import torch
import torch.nn as nn

class WatermarkRemover(nn.Module):
    def __init__(self):
        super(WatermarkRemover, self).__init__()
        self.enc1 = self.conv_block(3, 64)
        self.enc2 = self.conv_block(64, 128)
        self.enc3 = self.conv_block(128, 256)
        self.enc4 = self.conv_block(256, 512)

        self.bottleneck = self.conv_block(512, 1024)

        self.dec4 = self.conv_block(1024 + 512, 512)
        self.dec3 = self.conv_block(512 + 256, 256)
        self.dec2 = self.conv_block(256 + 128, 128)
        self.dec1 = self.conv_block(128 + 64, 64)

        self.final_layer = nn.Sequential(
            nn.Conv2d(64, 3, kernel_size=1),
            nn.Sigmoid()
        )

    def conv_block(self, in_channels, out_channels):
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )
    
    def UpsampleAndPadding(self, small_tensor, big_tensor):
        small_tensor = nn.Upsample(scale_factor=2)(small_tensor)

        b, c, ph, pw = small_tensor.shape
        h, w = big_tensor.shape[2:]
        padding_tensor = torch.zeros(b, c, h, w, device=big_tensor.device)
        padding_tensor[:, :, :ph, :pw] = small_tensor
        return padding_tensor


    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(nn.MaxPool2d(2)(e1))
        e3 = self.enc3(nn.MaxPool2d(2)(e2))
        e4 = self.enc4(nn.MaxPool2d(2)(e3))

        b = self.bottleneck(nn.MaxPool2d(2)(e4))

        d4 = self.dec4(torch.cat((self.UpsampleAndPadding(b, e4), e4), dim=1))
        d3 = self.dec3(torch.cat((self.UpsampleAndPadding(d4, e3), e3), dim=1))
        d2 = self.dec2(torch.cat((self.UpsampleAndPadding(d3, e2), e2), dim=1))
        d1 = self.dec1(torch.cat((self.UpsampleAndPadding(d2, e1), e1), dim=1))

        return self.final_layer(d1)

if __name__ == '__main__':
    model = WatermarkRemover()
    inp = torch.rand(2, 3, 240, 254)
    out = model(inp)
    print(out.shape)